import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from configs import DEVICE, EPOCHS, LEARNING_RATE, TEST_BATCH_SIZE, IMG_SIZE, IMG_CHANNEL, PATCH_SIZE, EMB_DIM, HEAD_NUM, MLP_RATIO, DEPTH, CLASS_NUM, HEAD_DIM
from model import ViTPred, MyViTPred
from dataloader import prepare_test_data

def inference_model(model,device,test_loader):
    model.eval()
    model.load_state_dict(torch.load('./ckpt/ViT'+str(EPOCHS)+'.pth',map_location=torch.device(DEVICE)))
    correct = 0
    test_loss = 0
    with torch.no_grad():
        for data,label in test_loader:
            data, label = data.to(device), label.to(device)
            output = model(data)
            #test_loss += F.cross_entropy(output,label).item()
            #predict = torch.max(torch.tensor(output),dim=1).indices
            #correct += predict.eq(label.view_as(predict)).sum().item()
            # print("Label of first data is:")
            # print(label[0])
            # print("Output of first data is:")
            print(output[0])
            # print("Shape of label of first batch is:")
            # print(label.shape)
            # print("Shape of output of first batch is:")
            # print(output.shape)
            # print("Len of output of first data is:")
            # print(len(output[0]))
            print(torch.isnan(label).any())
            assert (not torch.isnan(output).any()), "输出中存在nan"
            break
        # test_loss /= len(test_loader.dataset)
        # correct /= len(test_loader.dataset)
        # print("Test--Average loss:{:.4f},Accuracy:{:.3f}\n".format(test_loss,100.0*correct))

def inference():
    
    #model = ViTPred(img_size=IMG_SIZE, img_channel=IMG_CHANNEL, patch_size=PATCH_SIZE, emb_dim=EMB_DIM, batch_size=TEST_BATCH_SIZE, head_num=HEAD_NUM, mlp_ratio=MLP_RATIO, depth=DEPTH,class_num=CLASS_NUM).to(DEVICE)
    model = MyViTPred(img_size=IMG_SIZE, img_channel=IMG_CHANNEL, patch_size=PATCH_SIZE, emb_dim=EMB_DIM, batch_size=TEST_BATCH_SIZE, head_num=HEAD_NUM, mlp_ratio=MLP_RATIO, depth=DEPTH,class_num=CLASS_NUM,head_dim=HEAD_DIM).to(DEVICE)

    test_loader = prepare_test_data()
    inference_model(model,DEVICE,test_loader)

inference()